"""
AI Agent Pro - 记忆管理系统
"""

import asyncio
import json
import time
import uuid
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional, Union
from enum import Enum

import numpy as np
from sqlalchemy import Column, String, DateTime, Text, Float, Integer, Boolean
from sqlalchemy.ext.asyncio import AsyncSession
from pydantic import BaseModel, Field

from app.core.config import get_settings
from app.core.database import Base, get_db, get_qdrant
from app.core.logging import get_logger
from app.core.exceptions import MemoryError

logger = get_logger(__name__)
settings = get_settings()


class MemoryType(str, Enum):
    """记忆类型枚举"""
    CONVERSATION = "conversation"    # 对话记忆
    KNOWLEDGE = "knowledge"          # 知识记忆
    EXPERIENCE = "experience"        # 经验记忆
    FACT = "fact"                   # 事实记忆
    SKILL = "skill"                 # 技能记忆
    CONTEXT = "context"             # 上下文记忆


class MemoryImportance(str, Enum):
    """记忆重要性枚举"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


class Memory(Base):
    """记忆数据库模型"""
    __tablename__ = "memories"
    
    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    agent_id = Column(String, nullable=False, index=True)
    user_id = Column(String, nullable=True, index=True)
    memory_type = Column(String, nullable=False)
    content = Column(Text, nullable=False)
    summary = Column(Text, nullable=True)
    importance = Column(String, default=MemoryImportance.MEDIUM.value)
    access_count = Column(Integer, default=0)
    last_accessed = Column(DateTime, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    metadata = Column(Text, nullable=True)  # JSON存储额外信息
    embedding_id = Column(String, nullable=True)  # 向量数据库中的ID
    is_deleted = Column(Boolean, default=False)


class MemorySchema(BaseModel):
    """记忆数据模式"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    agent_id: str
    user_id: Optional[str] = None
    memory_type: MemoryType
    content: str
    summary: Optional[str] = None
    importance: MemoryImportance = MemoryImportance.MEDIUM
    access_count: int = 0
    last_accessed: Optional[datetime] = None
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    metadata: Dict[str, Any] = Field(default_factory=dict)
    embedding_id: Optional[str] = None
    similarity_score: Optional[float] = None  # 搜索时的相似度分数


class MemorySearchResult(BaseModel):
    """记忆搜索结果"""
    memories: List[MemorySchema]
    total_count: int
    search_time: float
    query_embedding: Optional[List[float]] = None


class MemoryManager:
    """记忆管理器"""
    
    def __init__(self):
        self.embedding_model = None
        self.vector_client = None
        self.collection_name = "agent_memory"
        
    async def initialize(self):
        """初始化记忆管理器"""
        try:
            # 初始化向量数据库客户端
            self.vector_client = await get_qdrant()
            
            # 初始化嵌入模型 (这里使用OpenAI的嵌入模型)
            if settings.openai_api_key:
                from openai import AsyncOpenAI
                self.embedding_client = AsyncOpenAI(api_key=settings.openai_api_key)
            
            logger.info("记忆管理器初始化成功")
            
        except Exception as e:
            logger.error(f"记忆管理器初始化失败: {str(e)}")
            raise MemoryError("初始化", str(e))
    
    async def save_memory(
        self,
        agent_id: str,
        content: str,
        memory_type: MemoryType = MemoryType.CONVERSATION,
        user_id: Optional[str] = None,
        importance: MemoryImportance = MemoryImportance.MEDIUM,
        metadata: Optional[Dict[str, Any]] = None
    ) -> MemorySchema:
        """
        保存记忆
        
        Args:
            agent_id: 智能体ID
            content: 记忆内容
            memory_type: 记忆类型
            user_id: 用户ID
            importance: 重要性
            metadata: 元数据
            
        Returns:
            保存的记忆对象
        """
        try:
            # 创建记忆对象
            memory = MemorySchema(
                agent_id=agent_id,
                user_id=user_id,
                memory_type=memory_type,
                content=content,
                importance=importance,
                metadata=metadata or {}
            )
            
            # 生成摘要
            memory.summary = await self._generate_summary(content)
            
            # 生成嵌入向量
            embedding = await self._create_embedding(content)
            
            # 保存到向量数据库
            if embedding and self.vector_client:
                embedding_id = await self._save_to_vector_db(memory, embedding)
                memory.embedding_id = embedding_id
            
            # 保存到关系数据库
            async with AsyncSession() as session:
                db_memory = Memory(
                    id=memory.id,
                    agent_id=memory.agent_id,
                    user_id=memory.user_id,
                    memory_type=memory.memory_type.value,
                    content=memory.content,
                    summary=memory.summary,
                    importance=memory.importance.value,
                    metadata=json.dumps(memory.metadata),
                    embedding_id=memory.embedding_id
                )
                
                session.add(db_memory)
                await session.commit()
            
            logger.info(f"记忆保存成功: {memory.id}")
            return memory
            
        except Exception as e:
            logger.error(f"保存记忆失败: {str(e)}")
            raise MemoryError("保存", str(e))
    
    async def search_memories(
        self,
        query: str,
        agent_id: Optional[str] = None,
        user_id: Optional[str] = None,
        memory_types: Optional[List[MemoryType]] = None,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        time_range: Optional[tuple] = None
    ) -> MemorySearchResult:
        """
        搜索记忆
        
        Args:
            query: 搜索查询
            agent_id: 智能体ID过滤
            user_id: 用户ID过滤
            memory_types: 记忆类型过滤
            limit: 返回数量限制
            similarity_threshold: 相似度阈值
            time_range: 时间范围过滤 (start_time, end_time)
            
        Returns:
            搜索结果
        """
        start_time = time.time()
        
        try:
            # 生成查询嵌入
            query_embedding = await self._create_embedding(query)
            
            if not query_embedding or not self.vector_client:
                # 回退到文本搜索
                return await self._fallback_text_search(
                    query, agent_id, user_id, memory_types, limit, time_range
                )
            
            # 构建过滤条件
            search_filter = {}
            if agent_id:
                search_filter["agent_id"] = agent_id
            if user_id:
                search_filter["user_id"] = user_id
            if memory_types:
                search_filter["memory_type"] = [t.value for t in memory_types]
            
            # 向量搜索
            search_results = await self._vector_search(
                query_embedding,
                search_filter,
                limit,
                similarity_threshold
            )
            
            # 获取详细记忆信息
            memories = []
            for result in search_results:
                memory = await self._get_memory_by_embedding_id(result["id"])
                if memory:
                    memory.similarity_score = result["score"]
                    
                    # 时间范围过滤
                    if time_range:
                        start, end = time_range
                        if not (start <= memory.created_at <= end):
                            continue
                    
                    memories.append(memory)
            
            search_time = time.time() - start_time
            
            return MemorySearchResult(
                memories=memories,
                total_count=len(memories),
                search_time=search_time,
                query_embedding=query_embedding[:10]  # 只返回前10个维度作为示例
            )
            
        except Exception as e:
            logger.error(f"搜索记忆失败: {str(e)}")
            raise MemoryError("搜索", str(e))
    
    async def get_memory(self, memory_id: str) -> Optional[MemorySchema]:
        """获取特定记忆"""
        try:
            async with AsyncSession() as session:
                result = await session.get(Memory, memory_id)
                
                if result and not result.is_deleted:
                    # 更新访问统计
                    result.access_count += 1
                    result.last_accessed = datetime.utcnow()
                    await session.commit()
                    
                    return self._db_to_schema(result)
                
                return None
                
        except Exception as e:
            logger.error(f"获取记忆失败: {str(e)}")
            raise MemoryError("获取", str(e))
    
    async def update_memory(
        self,
        memory_id: str,
        content: Optional[str] = None,
        importance: Optional[MemoryImportance] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> Optional[MemorySchema]:
        """更新记忆"""
        try:
            async with AsyncSession() as session:
                memory = await session.get(Memory, memory_id)
                
                if not memory or memory.is_deleted:
                    return None
                
                # 更新字段
                if content is not None:
                    memory.content = content
                    memory.summary = await self._generate_summary(content)
                    
                    # 更新嵌入向量
                    if memory.embedding_id:
                        embedding = await self._create_embedding(content)
                        if embedding:
                            await self._update_vector_db(memory.embedding_id, embedding)
                
                if importance is not None:
                    memory.importance = importance.value
                
                if metadata is not None:
                    memory.metadata = json.dumps(metadata)
                
                memory.updated_at = datetime.utcnow()
                await session.commit()
                
                return self._db_to_schema(memory)
                
        except Exception as e:
            logger.error(f"更新记忆失败: {str(e)}")
            raise MemoryError("更新", str(e))
    
    async def delete_memory(self, memory_id: str, soft_delete: bool = True) -> bool:
        """删除记忆"""
        try:
            async with AsyncSession() as session:
                memory = await session.get(Memory, memory_id)
                
                if not memory:
                    return False
                
                if soft_delete:
                    # 软删除
                    memory.is_deleted = True
                    memory.updated_at = datetime.utcnow()
                else:
                    # 硬删除
                    if memory.embedding_id:
                        await self._delete_from_vector_db(memory.embedding_id)
                    await session.delete(memory)
                
                await session.commit()
                return True
                
        except Exception as e:
            logger.error(f"删除记忆失败: {str(e)}")
            raise MemoryError("删除", str(e))
    
    async def save_conversation(
        self,
        agent_id: str,
        user_id: Optional[str],
        user_message: str,
        assistant_response: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> MemorySchema:
        """保存对话到记忆"""
        conversation_content = f"用户: {user_message}\n助手: {assistant_response}"
        
        conversation_metadata = {
            "type": "conversation",
            "user_message": user_message,
            "assistant_response": assistant_response,
            "timestamp": datetime.utcnow().isoformat()
        }
        
        if metadata:
            conversation_metadata.update(metadata)
        
        return await self.save_memory(
            agent_id=agent_id,
            content=conversation_content,
            memory_type=MemoryType.CONVERSATION,
            user_id=user_id,
            importance=MemoryImportance.MEDIUM,
            metadata=conversation_metadata
        )
    
    async def get_recent_memories(
        self,
        agent_id: str,
        user_id: Optional[str] = None,
        limit: int = 10,
        memory_types: Optional[List[MemoryType]] = None
    ) -> List[MemorySchema]:
        """获取最近的记忆"""
        try:
            # 实现最近记忆获取逻辑
            # 这里简化实现，实际应该从数据库查询
            recent_memories = []
            
            # 模拟返回一些记忆
            for i in range(min(limit, 3)):
                memory = MemorySchema(
                    agent_id=agent_id,
                    user_id=user_id,
                    memory_type=MemoryType.CONVERSATION,
                    content=f"最近的对话记忆 {i+1}",
                    summary=f"对话摘要 {i+1}",
                    importance=MemoryImportance.MEDIUM,
                    created_at=datetime.utcnow() - timedelta(hours=i)
                )
                recent_memories.append(memory)
            
            return recent_memories
            
        except Exception as e:
            logger.error(f"获取最近记忆失败: {str(e)}")
            return []
    
    async def cleanup_old_memories(
        self,
        agent_id: str,
        retention_days: int = None
    ) -> int:
        """清理过期记忆"""
        if retention_days is None:
            retention_days = settings.memory_retention_days
        
        try:
            cutoff_date = datetime.utcnow() - timedelta(days=retention_days)
            
            # 这里应该实现真实的清理逻辑
            # 删除过期的低重要性记忆
            cleaned_count = 0
            
            logger.info(f"清理了 {cleaned_count} 条过期记忆")
            return cleaned_count
            
        except Exception as e:
            logger.error(f"清理记忆失败: {str(e)}")
            return 0
    
    async def get_memory_statistics(
        self,
        agent_id: str,
        user_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """获取记忆统计信息"""
        try:
            # 实现统计逻辑
            stats = {
                "total_memories": 0,
                "by_type": {},
                "by_importance": {},
                "recent_access_count": 0,
                "storage_size_mb": 0.0,
                "oldest_memory": None,
                "newest_memory": None
            }
            
            # 模拟统计数据
            stats.update({
                "total_memories": 156,
                "by_type": {
                    "conversation": 120,
                    "knowledge": 25,
                    "experience": 8,
                    "fact": 3
                },
                "by_importance": {
                    "low": 89,
                    "medium": 52,
                    "high": 12,
                    "critical": 3
                },
                "recent_access_count": 45,
                "storage_size_mb": 12.5
            })
            
            return stats
            
        except Exception as e:
            logger.error(f"获取记忆统计失败: {str(e)}")
            return {}
    
    # 私有方法
    
    async def _create_embedding(self, text: str) -> Optional[List[float]]:
        """创建文本嵌入向量"""
        try:
            if not self.embedding_client:
                # 模拟嵌入向量
                import random
                return [random.random() for _ in range(1536)]
            
            response = await self.embedding_client.embeddings.create(
                model="text-embedding-ada-002",
                input=text
            )
            
            return response.data[0].embedding
            
        except Exception as e:
            logger.warning(f"创建嵌入向量失败: {str(e)}")
            return None
    
    async def _generate_summary(self, content: str, max_length: int = 100) -> str:
        """生成内容摘要"""
        if len(content) <= max_length:
            return content
        
        # 简单截取，实际应该使用AI生成摘要
        return content[:max_length] + "..."
    
    async def _save_to_vector_db(
        self,
        memory: MemorySchema,
        embedding: List[float]
    ) -> Optional[str]:
        """保存到向量数据库"""
        try:
            if not self.vector_client:
                return None
            
            # 生成向量ID
            vector_id = str(uuid.uuid4())
            
            # 保存到Qdrant (这里简化实现)
            # 实际应该调用Qdrant的API
            logger.info(f"向量保存到数据库: {vector_id}")
            
            return vector_id
            
        except Exception as e:
            logger.warning(f"保存到向量数据库失败: {str(e)}")
            return None
    
    async def _vector_search(
        self,
        query_embedding: List[float],
        filters: Dict[str, Any],
        limit: int,
        threshold: float
    ) -> List[Dict[str, Any]]:
        """向量搜索"""
        try:
            # 模拟向量搜索结果
            mock_results = []
            for i in range(min(limit, 5)):
                mock_results.append({
                    "id": f"vector_{i}",
                    "score": 0.9 - i * 0.1,
                    "payload": filters
                })
            
            return mock_results
            
        except Exception as e:
            logger.warning(f"向量搜索失败: {str(e)}")
            return []
    
    async def _fallback_text_search(
        self,
        query: str,
        agent_id: Optional[str],
        user_id: Optional[str],
        memory_types: Optional[List[MemoryType]],
        limit: int,
        time_range: Optional[tuple]
    ) -> MemorySearchResult:
        """回退文本搜索"""
        # 简化实现：模拟文本搜索
        memories = []
        
        # 模拟搜索结果
        for i in range(min(limit, 3)):
            memory = MemorySchema(
                agent_id=agent_id or "default",
                user_id=user_id,
                memory_type=MemoryType.CONVERSATION,
                content=f"包含'{query}'的记忆内容 {i+1}",
                summary=f"关于{query}的摘要",
                importance=MemoryImportance.MEDIUM,
                similarity_score=0.8 - i * 0.1
            )
            memories.append(memory)
        
        return MemorySearchResult(
            memories=memories,
            total_count=len(memories),
            search_time=0.1
        )
    
    async def _get_memory_by_embedding_id(self, embedding_id: str) -> Optional[MemorySchema]:
        """通过嵌入ID获取记忆"""
        # 简化实现
        return MemorySchema(
            agent_id="default",
            memory_type=MemoryType.CONVERSATION,
            content="通过嵌入ID获取的记忆",
            summary="嵌入记忆摘要",
            embedding_id=embedding_id
        )
    
    async def _update_vector_db(self, embedding_id: str, embedding: List[float]):
        """更新向量数据库"""
        logger.info(f"更新向量数据库: {embedding_id}")
    
    async def _delete_from_vector_db(self, embedding_id: str):
        """从向量数据库删除"""
        logger.info(f"从向量数据库删除: {embedding_id}")
    
    def _db_to_schema(self, db_memory: Memory) -> MemorySchema:
        """数据库模型转换为Pydantic模型"""
        metadata = {}
        if db_memory.metadata:
            try:
                metadata = json.loads(db_memory.metadata)
            except json.JSONDecodeError:
                pass
        
        return MemorySchema(
            id=db_memory.id,
            agent_id=db_memory.agent_id,
            user_id=db_memory.user_id,
            memory_type=MemoryType(db_memory.memory_type),
            content=db_memory.content,
            summary=db_memory.summary,
            importance=MemoryImportance(db_memory.importance),
            access_count=db_memory.access_count,
            last_accessed=db_memory.last_accessed,
            created_at=db_memory.created_at,
            updated_at=db_memory.updated_at,
            metadata=metadata,
            embedding_id=db_memory.embedding_id
        )


# 全局记忆管理器实例
memory_manager = MemoryManager()


# 导出
__all__ = [
    "MemoryManager",
    "MemorySchema",
    "MemoryType",
    "MemoryImportance",
    "MemorySearchResult",
    "Memory",
    "memory_manager",
]
